#!/usr/bin/env python3
"""
颜色校准工具 - 保留用于传统颜色检测方法
"""
import cv2
import numpy as np


def nothing(x):
    pass


def calibrate_color():
    # 创建窗口和轨迹栏
    cv2.namedWindow('Calibration')

    # 红色需要两组HSV范围
    cv2.createTrackbar('H1 Lower', 'Calibration', 0, 180, nothing)
    cv2.createTrackbar('H1 Upper', 'Calibration', 10, 180, nothing)
    cv2.createTrackbar('H2 Lower', 'Calibration', 170, 180, nothing)
    cv2.createTrackbar('H2 Upper', 'Calibration', 180, 180, nothing)

    cv2.createTrackbar('S Lower', 'Calibration', 120, 255, nothing)
    cv2.createTrackbar('S Upper', 'Calibration', 255, 255, nothing)
    cv2.createTrackbar('V Lower', 'Calibration', 70, 255, nothing)
    cv2.createTrackbar('V Upper', 'Calibration', 255, 255, nothing)

    # 初始化摄像头
    cap = cv2.VideoCapture(0)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

    print("使用轨迹栏调整HSV范围，按's'保存当前值，'q'退出")

    try:
        while True:
            ret, frame = cap.read()
            if not ret:
                break

            # 获取轨迹栏位置
            h1_l = cv2.getTrackbarPos('H1 Lower', 'Calibration')
            h1_u = cv2.getTrackbarPos('H1 Upper', 'Calibration')
            h2_l = cv2.getTrackbarPos('H2 Lower', 'Calibration')
            h2_u = cv2.getTrackbarPos('H2 Upper', 'Calibration')
            s_l = cv2.getTrackbarPos('S Lower', 'Calibration')
            s_u = cv2.getTrackbarPos('S Upper', 'Calibration')
            v_l = cv2.getTrackbarPos('V Lower', 'Calibration')
            v_u = cv2.getTrackbarPos('V Upper', 'Calibration')

            # 转换到HSV并创建掩膜
            hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)

            lower_red1 = np.array([h1_l, s_l, v_l])
            upper_red1 = np.array([h1_u, s_u, v_u])
            lower_red2 = np.array([h2_l, s_l, v_l])
            upper_red2 = np.array([h2_u, s_u, v_u])

            mask1 = cv2.inRange(hsv, lower_red1, upper_red1)
            mask2 = cv2.inRange(hsv, lower_red2, upper_red2)
            mask = mask1 + mask2

            # 显示原图和掩膜
            result = cv2.bitwise_and(frame, frame, mask=mask)
            cv2.imshow('Original', frame)
            cv2.imshow('Mask', mask)
            cv2.imshow('Result', result)

            key = cv2.waitKey(1) & 0xFF
            if key == ord('q'):
                break
            elif key == ord('s'):
                # 保存HSV范围
                ranges = {
                    'lower1': [h1_l, s_l, v_l],
                    'upper1': [h1_u, s_u, v_u],
                    'lower2': [h2_l, s_l, v_l],
                    'upper2': [h2_u, s_u, v_u]
                }
                print(f"当前HSV范围: {ranges}")
                print("将这些值更新到配置文件中")

    finally:
        cap.release()
        cv2.destroyAllWindows()


if __name__ == "__main__":
    calibrate_color()